import numpy

def find_k(img, z):
    maximum_intensity = img[:, :, z].flatten().max()
    #print maximum_intensity
    
    prob, bin_edges = numpy.histogram(img[:, :, z].flatten(), maximum_intensity + 1, normed=True)
    #prob /= maximum_intensity
    N = len(prob)
    #print N, prob
    
    mu_t = numpy.multiply(prob, [i for i in xrange(N)]).sum()
    
    omega_1 = prob[0] / maximum_intensity
    omega_2 = 1. - omega_1
    mu_1 = 0
    mu_2 = (mu_t - mu_1) / omega_2    
    max_val = 0.
    k = 0
    
    #print mu_t
    
    for t in xrange(int(N)):        
        omega_1 += prob[t] / maximum_intensity
        omega_2 = 1. - omega_1
        mu_1 += t * prob[t] 
        mu_2 = (mu_t - mu_1) / omega_2
        #print omega_1, omega_2
        #print (omega_1 * omega_2) * (mu_2 - mu_1) ** 2
        #print " "
        # Varianza intra clase
        sigma_b = (omega_1 * omega_2) * (mu_2 - mu_1) ** 2
        # Busco el maximo
        if sigma_b > max_val:
            max_val = sigma_b
            k = t
    
    return k * 1.
    
def segmentation(img, z):
    k = find_k(img, z)
    print k
    
    res = numpy.empty((img.shape[0], img.shape[1], img.shape[2]), dtype = float)
    
    for i in xrange(img.shape[0]):
        for j in xrange(img.shape[1]):
            for k in xrange(z, z + 1):
                #print img[i, j, k]
                res[i, j, k] = 0.3 if img[i, j, k] <= k else 0.7
    
    return res
